package io.kiki.sba.registry.server.metadata.slot.balance;

import io.kiki.sba.registry.common.model.Tuple;
import io.kiki.sba.registry.common.model.slot.DataStoreNode;
import io.kiki.sba.registry.common.model.slot.Slots;
import io.kiki.sba.registry.server.metadata.slot.SlotBalancer;
import io.kiki.sba.registry.server.metadata.slot.util.builder.SlotsBuilder;
import io.kiki.sba.registry.server.metadata.slot.util.comparator.Comparators;
import io.kiki.sba.registry.server.shared.slot.SlotsUtils;
import io.kiki.sba.registry.util.MathUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;


public class LeaderOnlyBalancer implements SlotBalancer {

    public static final int TRIGGER_THESHOLD = 2;
    private static final Logger logger = LoggerFactory.getLogger(LeaderOnlyBalancer.class);
    private final SlotsBuilder slotsBuilder;

    private final Set<String> currentDataStoreNodeSet;

    private final BalancePolicyImpl balancePolicy = new BalancePolicyImpl();

    public LeaderOnlyBalancer(SlotsBuilder slotsBuilder, Collection<String> currentDataStoreNodeSet) {
        this.slotsBuilder = slotsBuilder;
        this.currentDataStoreNodeSet = Collections.unmodifiableSet(new TreeSet<>(currentDataStoreNodeSet));
    }

    @Override
    public Slots balance() {
        final int leaderAvgNumPerDataStoreNode = MathUtils.divideCeil(slotsBuilder.getSlotNums(), currentDataStoreNodeSet.size());
        final int maxMoveLeaderSlots = balancePolicy.getMaxMoveLeaderSlots();
        int balanced = 0;
        while (balanced < maxMoveLeaderSlots) {
            List<String> targetHighDataStoreNodeList = findDataStoreNodeListWithLeaderHighWaterMark(leaderAvgNumPerDataStoreNode);
            if (targetHighDataStoreNodeList.isEmpty()) {
                logger.info("[balance]no data-servers needs leader balance");
                // no balance, return null
                break;
            }
            final Set<String> excludes = new HashSet<>(targetHighDataStoreNodeList);
            /**
             * int threshold = balancePolicy.getHighWaterMarkSlotLeaderNums(avg); >100%,所以需要做下面的操作.
             */
            // exclude the dataNode which could not add any leader
            excludes.addAll(findDataStoreNodeListWithLeaderHighWaterMark(leaderAvgNumPerDataStoreNode - 1));

            for (String target : targetHighDataStoreNodeList) {
                Tuple<String, Integer> tuple = selectLeader(target, excludes);
                if (tuple == null) {
                    continue;
                }
                final String newLeader = tuple.o1;
                final int slotId = tuple.o2;
                slotsBuilder.replaceLeaderDataStoreNodeId(slotId, newLeader);
                balanced++;
                break;
            }
        }
        if (balanced != 0) {
            Slots slots = slotsBuilder.build();
            logger.info("[balance][end] slot table leader stat: {}", SlotsUtils.getSlotsLeaderCount(slots));
            logger.info("[balance][end] slot table slots stat: {}", SlotsUtils.getSlotsSlotCount(slots));
            return slots;
        }
        return null;
    }

    private List<String> findDataStoreNodeListWithLeaderHighWaterMark(int avg) {
        int threshold = balancePolicy.getHighWaterMarkSlotLeaderNums(avg);
        List<DataStoreNode> dataStoreNodeList = slotsBuilder.getDataNodeSlotsLeaderBeyond(threshold);
        List<String> dataStoreNodeIdList = DataStoreNode.collectDataStoreNodeIdList(dataStoreNodeList);
        dataStoreNodeIdList.sort(Comparators.mostLeadersFirst(slotsBuilder));
        logger.info("[LeaderHighWaterMark] avg={}, threshold={}, dataServers={}", avg, threshold, dataStoreNodeIdList);
        return dataStoreNodeIdList;
    }

    private Tuple<String, Integer> selectLeader(String leaderDataStoreNode, Set<String> excludes) {
        final DataStoreNode dataStoreNode = slotsBuilder.getDataStoreNode(leaderDataStoreNode);
        /**
         *  Comparators.leastLeadersFirst(slotsBuilder)
         */
        List<String> candidates = getCandidateDataServers(currentDataStoreNodeSet, excludes, Comparators.leastLeadersFirst(slotsBuilder));
        /**
         * 为什么这里只取第一个? 需要确认!
         */
        for (int leaderSlotId : dataStoreNode.getLeaderSlotIdSet()) {
            for (String candidate : candidates) {
                return Tuple.of(candidate, leaderSlotId);
            }
        }
        return null;
    }

    private List<String> getCandidateDataServers(Collection<String> candidateDataServers, Collection<String> excludes, Comparator<String> comparator) {
        Set<String> candidates = new HashSet<>(candidateDataServers);
        candidates.removeAll(excludes);
        List<String> list = new ArrayList<>(candidates);
        if (comparator != null) {
            list.sort(comparator);
        }
        return list;
    }


}
